# ==============================================================================
# SCRIPT 03: CLUSTERIZAÇÃO K-MEANS (2022 - k=3)
# ==============================================================================

# ------------------------------------------------------------------------------
# 1. OBJETIVO
# ------------------------------------------------------------------------------
# Este script executa a clusterização K-Means com k=3 nos dados de 2022.
# Usamos k=3 (ao invés do k=5 ótimo) para evitar overfitting de clusters
# com apenas 2 RPPS.

# ------------------------------------------------------------------------------
# 2. BIBLIOTECAS NECESSÁRIAS
# ------------------------------------------------------------------------------
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
import warnings

warnings.filterwarnings("ignore")

# ------------------------------------------------------------------------------
# 3. CONFIGURAÇÕES E CONSTANTES
# ------------------------------------------------------------------------------

# Caminhos dos arquivos
INPUT_PATH = "../2022_data/base_final_2022.csv"
OUTPUT_PATH = "../2022_data/base_final_2022_com_clusters_k3.csv"

# Número de clusters
K = 3

# Segmentos de investimento
SEGMENTOS = [
    "Renda Fixa", "Renda Variável", "Investimentos no Exterior",
    "Investimentos Estruturados", "Fundos Imobiliários",
    "Empréstimos Consignados", "Disponibilidades Financeiras"
]

# ------------------------------------------------------------------------------
# 4. EXECUÇÃO DO SCRIPT
# ------------------------------------------------------------------------------

print("="*80)
print(f"INICIANDO CLUSTERIZAÇÃO K-MEANS (2022 - k={K})")
print("="*80)

# --- 4.1. Carregamento dos Dados ---
print("\n1. Carregando dados...")
df = pd.read_csv(INPUT_PATH)
print(f"  - {len(df)} RPPS carregados.")

# --- 4.2. Padronização dos Dados ---
print("\n2. Padronizando dados...")

# Extrair apenas os segmentos para clusterização
X = df[SEGMENTOS].values

# Padronizar os dados (Z-score)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
print(f"  - Dados padronizados com sucesso.")

# --- 4.3. Execução do K-Means ---
print(f"\n3. Executando K-Means com k={K}...")

# Configurar o K-Means com random_state fixo para reprodutibilidade
kmeans = KMeans(n_clusters=K, random_state=42, n_init=10, max_iter=300)

# Treinar o modelo
df["Cluster"] = kmeans.fit_predict(X_scaled)
print(f"  - Clusterização concluída.")

# --- 4.4. Análise dos Clusters ---
print(f"\n4. Analisando os clusters...")

for cluster_id in range(K):
    cluster_size = (df["Cluster"] == cluster_id).sum()
    cluster_pct = (cluster_size / len(df)) * 100
    print(f"  - Cluster {cluster_id}: {cluster_size} RPPS ({cluster_pct:.1f}%)")

# Calcular perfil médio de cada cluster
print("\n5. Perfil médio de cada cluster:")
for cluster_id in range(K):
    print(f"\n  Cluster {cluster_id}:")
    cluster_data = df[df["Cluster"] == cluster_id][SEGMENTOS]
    for segmento in SEGMENTOS:
        media = cluster_data[segmento].mean()
        print(f"    - {segmento}: {media:.2f}%")

# --- 4.5. Salvando os Resultados ---
print(f"\n6. Salvando resultados em: {OUTPUT_PATH}")
df.to_csv(OUTPUT_PATH, index=False)

print("\n" + "="*80)
print("CLUSTERIZAÇÃO DE 2022 CONCLUÍDA COM SUCESSO!")
print("="*80)
